package main

import (
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/signal"
	"sort"
	"sync"
	"syscall"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
)

// TCPFlow 存储单个TCP流的信息
type TCPFlow struct {
	SrcIP   string
	SrcPort uint16
	DstIP   string
	DstPort uint16

	// 连接状态
	StartTime time.Time
	LastSeen  time.Time
	Status    string // ESTABLISHED, SYN_SENT, etc.

	// 流量统计
	BytesSent       uint64
	BytesReceived   uint64
	PacketsSent     uint64
	PacketsReceived uint64

	// TCP指标
	RetransmittedPackets uint64
	DuplicateAcks        uint64
	LostPackets          uint64
	MinRTT               time.Duration
	MaxRTT               time.Duration
	AvgRTT               time.Duration
	RTTSamples           []time.Duration

	// 序列号跟踪
	LastSentSeq     uint32
	LastAckSeq      uint32
	LastReceivedSeq uint32
	LastSentTime    time.Time
}

// FlowKey 用于标识唯一的TCP流
type FlowKey struct {
	SrcIP   string
	SrcPort uint16
	DstIP   string
	DstPort uint16
}

// 全局状态
var (
	tcpFlows      = make(map[FlowKey]*TCPFlow)
	flowMutex     = &sync.Mutex{}
	iface         string
	snapshotLen   int32         = 1600
	promiscuous   bool          = false
	timeout       time.Duration = pcap.BlockForever
	printInterval time.Duration = 5 * time.Second
)

func init() {
	flag.StringVar(&iface, "iface", "", "网络接口名称 (例如: eth0, en0)")
	flag.DurationVar(&printInterval, "interval", 5*time.Second, "统计信息打印间隔")
	flag.Parse()

	if iface == "" {
		log.Fatal("必须指定网络接口 (-iface)")
	}
}

// 创建FlowKey，确保顺序一致（源IP:端口 < 目标IP:端口）
func getFlowKey(srcIP net.IP, srcPort uint16, dstIP net.IP, dstPort uint16) FlowKey {
	srcStr := fmt.Sprintf("%s:%d", srcIP, srcPort)
	dstStr := fmt.Sprintf("%s:%d", dstIP, dstPort)

	// 确保key的一致性，不管数据包方向
	if srcStr < dstStr {
		return FlowKey{
			SrcIP:   srcIP.String(),
			SrcPort: srcPort,
			DstIP:   dstIP.String(),
			DstPort: dstPort,
		}
	}
	return FlowKey{
		SrcIP:   dstIP.String(),
		SrcPort: dstPort,
		DstIP:   srcIP.String(),
		DstPort: srcPort,
	}
}

// 更新TCP流的RTT信息
func updateRTT(flow *TCPFlow, tcp *layers.TCP, packetTime time.Time) {
	if tcp.ACK && tcp.Ack > 0 && flow.LastSentSeq > 0 && tcp.Ack >= flow.LastSentSeq {
		// 计算RTT：当前时间 - 最后发送数据包的时间
		rtt := packetTime.Sub(flow.LastSentTime)
		if rtt > 0 {
			flow.RTTSamples = append(flow.RTTSamples, rtt)

			// 更新RTT统计
			if flow.MinRTT == 0 || rtt < flow.MinRTT {
				flow.MinRTT = rtt
			}
			if rtt > flow.MaxRTT {
				flow.MaxRTT = rtt
			}

			// 计算平均RTT
			total := time.Duration(0)
			for _, t := range flow.RTTSamples {
				total += t
			}
			flow.AvgRTT = total / time.Duration(len(flow.RTTSamples))
		}
	}
}

// 处理TCP数据包
func handleTCPPacket(packet gopacket.Packet) {
	ipLayer := packet.Layer(layers.LayerTypeIPv4)
	if ipLayer == nil {
		return
	}
	ip, _ := ipLayer.(*layers.IPv4)

	tcpLayer := packet.Layer(layers.LayerTypeTCP)
	if tcpLayer == nil {
		return
	}
	tcp, _ := tcpLayer.(*layers.TCP)

	// 获取流密钥
	key := getFlowKey(ip.SrcIP, uint16(tcp.SrcPort), ip.DstIP, uint16(tcp.DstPort))
	isOriginalDirection := (key.SrcIP == ip.SrcIP.String() && key.SrcPort == uint16(tcp.SrcPort))

	flowMutex.Lock()
	defer flowMutex.Unlock()

	// 检查流是否已存在，不存在则创建
	flow, exists := tcpFlows[key]
	if !exists {
		flow = &TCPFlow{
			SrcIP:     key.SrcIP,
			SrcPort:   key.SrcPort,
			DstIP:     key.DstIP,
			DstPort:   key.DstPort,
			StartTime: packet.Metadata().Timestamp,
			Status:    "NEW",
		}
		tcpFlows[key] = flow
	}

	// 更新最后活动时间
	flow.LastSeen = packet.Metadata().Timestamp

	// 更新TCP状态
	if tcp.SYN && !tcp.ACK {
		flow.Status = "SYN_SENT"
	} else if tcp.SYN && tcp.ACK {
		flow.Status = "SYN_RECEIVED"
	} else if tcp.ACK && flow.Status != "ESTABLISHED" {
		flow.Status = "ESTABLISHED"
	} else if tcp.FIN {
		flow.Status = "FIN_WAIT"
	} else if tcp.RST {
		flow.Status = "RESET"
	}

	// 更新流量统计
	packetSize := uint64(len(packet.Data()))
	if isOriginalDirection {
		flow.BytesSent += packetSize
		flow.PacketsSent++
		flow.LastSentSeq = tcp.Seq
		flow.LastSentTime = packet.Metadata().Timestamp
	} else {
		flow.BytesReceived += packetSize
		flow.PacketsReceived++
		flow.LastReceivedSeq = tcp.Seq
	}

	// 更新确认序列号
	if tcp.ACK {
		flow.LastAckSeq = tcp.Ack
	}

	// 检测重传和丢包
	if tcp.SYN {
		// SYN包不视为重传
	} else if isOriginalDirection {
		// 检查是否为重传 (简化检测：相同序列号但不是第一个包)
		if tcp.Seq <= flow.LastAckSeq && flow.PacketsSent > 1 {
			flow.RetransmittedPackets++
		}
	} else {
		// 检测重复ACK
		if tcp.ACK && tcp.Ack == flow.LastAckSeq && flow.PacketsReceived > 1 {
			flow.DuplicateAcks++
			// 重复ACK通常意味着丢包
			flow.LostPackets++
		}
	}

	// 更新RTT
	updateRTT(flow, tcp, packet.Metadata().Timestamp)
}

// 打印统计信息
func printStats() {
	flowMutex.Lock()
	defer flowMutex.Unlock()

	fmt.Printf("\n===== TCP监控统计 (%s) =====\n", time.Now().Format(time.RFC3339))
	fmt.Printf("总TCP连接数: %d\n", len(tcpFlows))
	fmt.Println("--------------------------------------------------------------------------")
	fmt.Printf("%-42s %-12s %-10s %-10s %-8s %-8s %-8s %-10s\n",
		"连接", "状态", "上行(KB)", "下行(KB)", "丢包数", "重传数", "重复ACK", "平均RTT")
	fmt.Println("--------------------------------------------------------------------------")

	// 按活动时间排序
	flows := make([]*TCPFlow, 0, len(tcpFlows))
	for _, flow := range tcpFlows {
		flows = append(flows, flow)
	}
	sort.Slice(flows, func(i, j int) bool {
		return flows[i].LastSeen.After(flows[j].LastSeen)
	})

	// 只显示前20个连接
	displayCount := 20
	if len(flows) < displayCount {
		displayCount = len(flows)
	}

	for i := 0; i < displayCount; i++ {
		flow := flows[i]
		connStr := fmt.Sprintf("%s:%d -> %s:%d",
			flow.SrcIP, flow.SrcPort, flow.DstIP, flow.DstPort)
		upKB := float64(flow.BytesSent) / 1024
		downKB := float64(flow.BytesReceived) / 1024

		rttStr := "-"
		if flow.AvgRTT > 0 {
			rttStr = flow.AvgRTT.String()
		}

		fmt.Printf("%-42s %-12s %-10.2f %-10.2f %-8d %-8d %-8d %-10s\n",
			connStr,
			flow.Status,
			upKB,
			downKB,
			flow.LostPackets,
			flow.RetransmittedPackets,
			flow.DuplicateAcks,
			rttStr,
		)
	}

	fmt.Println("--------------------------------------------------------------------------")
}

// 清理过期连接
func cleanupOldFlows() {
	flowMutex.Lock()
	defer flowMutex.Unlock()

	now := time.Now()
	// 清理超过5分钟没有活动的连接
	expirationTime := 5 * time.Minute

	for key, flow := range tcpFlows {
		if now.Sub(flow.LastSeen) > expirationTime {
			delete(tcpFlows, key)
		}
	}
}

func main() {
	// 打开网络设备
	handle, err := pcap.OpenLive(iface, snapshotLen, promiscuous, timeout)
	if err != nil {
		log.Fatalf("无法打开设备 %s: %v", iface, err)
	}
	defer handle.Close()

	// 设置BPF过滤器，只捕获TCP数据包
	if err := handle.SetBPFFilter("tcp"); err != nil {
		log.Fatalf("无法设置过滤器: %v", err)
	}

	// 启动统计打印goroutine
	go func() {
		ticker := time.NewTicker(printInterval)
		defer ticker.Stop()

		for range ticker.C {
			printStats()
			cleanupOldFlows()
		}
	}()

	// 捕获中断信号
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// 开始捕获数据包
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packetChan := packetSource.Packets()

	fmt.Printf("开始监控TCP连接 on %s...\n", iface)
	fmt.Printf("每 %v 打印一次统计信息\n", printInterval)
	fmt.Println("按Ctrl+C退出")

	for {
		select {
		case packet := <-packetChan:
			handleTCPPacket(packet)
		case <-sigChan:
			fmt.Println("\n收到退出信号，正在保存最后的统计信息...")
			printStats()
			fmt.Println("监控已停止")
			return
		}
	}
}
